# client_main.py
import pygame
import math
import sys
import struct
import pickle
from server_client import ClientComm


def main():
    print("Starting client...")
    client = ClientComm()
    role = client.get_connection()
    if role is None:
        print("Failed to connect")
        sys.exit(1)

    print(f"Connected as {role[0]} {role[1]}")

    # Pygame setup
    print("Initializing Pygame...")
    pygame.init()
    vision_size = 600
    screen = pygame.display.set_mode((vision_size, vision_size + 100))  # 调整窗口大小
    pygame.display.set_caption(f"RM Simulator - {role[1]} {role[0]}")
    clock = pygame.time.Clock()
    font = pygame.font.Font(None, 20)

    print("Entering main loop...")

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # Get observation
        print("Getting observation...")
        vision, reward, done, info = client.get_observation()
        if vision is None:
            running = False
            break
        if done:
            running = False
            print("Episode done")
            break
        # clear screen
        screen.fill((0, 0, 0))
        # Display vision
        print("Rendering vision...")
        if vision is not None:
            # vision (h,w,3), swap for pygame (w,h,3)? surfarray.make_surface expects (w,h,3)
            vision_surf = pygame.surfarray.make_surface(vision.swapaxes(0, 1))
            screen.blit(vision_surf, (0, 0))

            # 在client_main.py中更新HUD部分代码如下：

        # HUD
        print("Updating HUD...")
        hud_y = vision_size + 10
        # 提取信息
        remaining_time = info.get("remaining_time", 0)
        red_score = info.get("red_score", 0)
        blue_score = info.get("blue_score", 0)
        own_hp = info.get("own_hp", 0)
        own_max_hp = info.get("own_max_hp", 1000)
        bullet = info.get("bullet", 0)
        max_bullet = info.get("max_bullet", 0)
        attack_cd = info.get("attack_cd", 0)
        reborn_cd = info.get("reborn_cd", 0)

        robot_team = role[0]  # 队伍信息
        robot_type = role[1]  # 类型信息

        # 组织信息为多列显示
        hud_lines = [
            [
                f"Robot: {robot_team.capitalize()} {robot_type.capitalize()}",
                f"Score - Red: {red_score:.0f} Blue: {blue_score:.0f}",
                f"Time: {remaining_time:.1f}s",
            ],
            [
                f"HP: {own_hp:.0f}/{own_max_hp:.0f}",
                f"Bullet: {bullet:.0f}/{max_bullet:.0f}",
                f"Attack CD: {attack_cd:.1f}s",
            ],
        ]

        # 如果机器人死亡，显示重生时间
        if reborn_cd > 0:
            hud_lines.append([f"Reborn in: {reborn_cd:.1f}s", "", ""])

        # 添加奖励信息（如果有）
        if reward != 0:
            hud_lines.append([f"Last Reward: {reward:.1f}", "", ""])

        # 渲染多列信息
        for i, line_items in enumerate(hud_lines):
            for j, text in enumerate(line_items):
                if text:  # 只渲染非空文本
                    text_surf = font.render(text, True, (255, 255, 255))
                    screen.blit(text_surf, (10 + j * 200, hud_y + i * 25))

        pygame.display.flip()

        # Get action from keyboard
        print("Getting action from keyboard...")
        keys = pygame.key.get_pressed()
        vel = 0.0
        rot_val = 0.0  # renamed to avoid conflict with module
        attack = 0

        if keys[pygame.K_w]:
            vel = 1.0
        # s: stop
        if keys[pygame.K_s]:
            vel = 0.0

        if keys[pygame.K_a]:
            rot_val = -1.0
        if keys[pygame.K_d]:
            rot_val = 1.0

        if keys[pygame.K_SPACE]:
            attack = 1.0

        # For combinations like w+a: vel=1, rot=-1 (curves to 45 deg approx)
        action = [vel, rot_val, attack]
        client.send_action(action)

        clock.tick(60)

    client.close()
    pygame.quit()


if __name__ == "__main__":
    main()
